#!/bin/bash

# 数据备份和恢复脚本
# 用于生产环境数据备份和恢复

set -e

# 配置
BACKUP_DIR="${BACKUP_DIR:-./backups}"
DATA_DIR="${DATA_DIR:-./data}"
LOG_DIR="${LOG_DIR:-./logs}"
TEMP_DIR="${TEMP_DIR:-./temp}"
CONFIG_DIR="${CONFIG_DIR:-./config}"

# 备份配置
BACKUP_RETENTION_DAYS="${BACKUP_RETENTION_DAYS:-30}"
BACKUP_COMPRESS="${BACKUP_COMPRESS:-true}"
BACKUP_ENCRYPT="${BACKUP_ENCRYPT:-false}"
ENCRYPTION_KEY="${ENCRYPTION_KEY:-}"
BACKUP_SCHEDULE="${BACKUP_SCHEDULE:-0 2 * * *}"

# 数据库配置
DATABASE_FILE="${DATABASE_FILE:-$DATA_DIR/app.db}"
DATABASE_URL="${DATABASE_URL:-sqlite+aiosqlite://$DATABASE_FILE}"

# Redis配置
REDIS_URL="${REDIS_URL:-redis://localhost:6379/0}"

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1"
}

# 检查命令是否存在
check_command() {
    if ! command -v $1 &> /dev/null; then
        log_error "Command '$1' not found"
        exit 1
    fi
}

# 创建备份目录
create_backup_dir() {
    local timestamp=$(date +%Y%m%d_%H%M%S)
    BACKUP_PATH="$BACKUP_DIR/backup_$timestamp"
    
    mkdir -p "$BACKUP_PATH"
    mkdir -p "$BACKUP_PATH/data"
    mkdir -p "$BACKUP_PATH/logs"
    mkdir -p "$BACKUP_PATH/config"
    mkdir -p "$BACKUP_PATH/database"
    mkdir -p "$BACKUP_PATH/temp"
    
    echo "$BACKUP_PATH"
}

# 备份数据库
backup_database() {
    local backup_path="$1"
    log_info "Backing up database..."
    
    if [ -f "$DATABASE_FILE" ]; then
        # SQLite数据库备份
        local db_backup="$backup_path/database/app.db"
        
        # 使用SQLite备份命令
        sqlite3 "$DATABASE_FILE" ".backup $db_backup" || {
            log_error "Failed to backup SQLite database"
            return 1
        }
        
        # 验证备份
        if [ -f "$db_backup" ]; then
            log_info "Database backup created: $db_backup"
            log_info "Database backup size: $(du -h "$db_backup" | cut -f1)"
        else
            log_error "Database backup file not found"
            return 1
        fi
    else
        log_warn "Database file not found: $DATABASE_FILE"
    fi
    
    return 0
}

# 备份Redis数据
backup_redis() {
    local backup_path="$1"
    log_info "Backing up Redis data..."
    
    # 解析Redis连接信息
    if [[ $REDIS_URL =~ redis://([^:]+):?([0-9]*)/([0-9]+) ]]; then
        local redis_host="${BASH_REMATCH[1]}"
        local redis_port="${BASH_REMATCH[2]}"
        local redis_db="${BASH_REMATCH[3]}"
        
        if [ -z "$redis_port" ]; then
            redis_port="6379"
        fi
        
        local redis_backup="$backup_path/database/redis.rdb"
        
        # 使用redis-cli备份数据
        redis-cli -h "$redis_host" -p "$redis_port" --rdb "$redis_backup" || {
            log_warn "Failed to backup Redis data, continuing..."
        }
        
        if [ -f "$redis_backup" ]; then
            log_info "Redis backup created: $redis_backup"
            log_info "Redis backup size: $(du -h "$redis_backup" | cut -f1)"
        else
            log_warn "Redis backup not created"
        fi
    else
        log_warn "Invalid Redis URL: $REDIS_URL"
    fi
    
    return 0
}

# 备份数据目录
backup_data() {
    local backup_path="$1"
    log_info "Backing up data directory..."
    
    if [ -d "$DATA_DIR" ]; then
        # 排除数据库文件（已单独备份）
        tar --exclude="*.db*" -cf "$backup_path/data/data_files.tar" -C "$DATA_DIR" . 2>/dev/null || {
            log_warn "Failed to backup some data files"
        }
        
        if [ -f "$backup_path/data/data_files.tar" ]; then
            log_info "Data files backup created"
            log_info "Data files backup size: $(du -h "$backup_path/data/data_files.tar" | cut -f1)"
        fi
    else
        log_warn "Data directory not found: $DATA_DIR"
    fi
    
    return 0
}

# 备份日志目录
backup_logs() {
    local backup_path="$1"
    log_info "Backing up logs directory..."
    
    if [ -d "$LOG_DIR" ]; then
        tar -cf "$backup_path/logs/logs.tar" -C "$LOG_DIR" . 2>/dev/null || {
            log_warn "Failed to backup logs"
        }
        
        if [ -f "$backup_path/logs/logs.tar" ]; then
            log_info "Logs backup created"
            log_info "Logs backup size: $(du -h "$backup_path/logs/logs.tar" | cut -f1)"
        fi
    else
        log_warn "Logs directory not found: $LOG_DIR"
    fi
    
    return 0
}

# 备份配置文件
backup_config() {
    local backup_path="$1"
    log_info "Backing up configuration files..."
    
    if [ -d "$CONFIG_DIR" ]; then
        tar -cf "$backup_path/config/config.tar" -C "$CONFIG_DIR" . 2>/dev/null || {
            log_warn "Failed to backup configuration"
        }
        
        if [ -f "$backup_path/config/config.tar" ]; then
            log_info "Configuration backup created"
            log_info "Configuration backup size: $(du -h "$backup_path/config/config.tar" | cut -f1)"
        fi
    else
        log_warn "Config directory not found: $CONFIG_DIR"
    fi
    
    # 备份环境变量文件
    if [ -f ".env" ]; then
        cp ".env" "$backup_path/config/.env"
        log_info "Environment file backed up"
    fi
    
    return 0
}

# 备份临时文件（选择性）
backup_temp() {
    local backup_path="$1"
    log_info "Backing up temporary files..."
    
    if [ -d "$TEMP_DIR" ]; then
        # 只备份特定类型的临时文件
        find "$TEMP_DIR" -type f \( -name "*.json" -o -name "*.xml" -o -name "*.yml" -o -name "*.yaml" \) | \
        tar -cf "$backup_path/temp/temp_files.tar" -T - 2>/dev/null || {
            log_warn "Failed to backup temporary files"
        }
        
        if [ -f "$backup_path/temp/temp_files.tar" ]; then
            log_info "Temporary files backup created"
            log_info "Temporary files backup size: $(du -h "$backup_path/temp/temp_files.tar" | cut -f1)"
        fi
    else
        log_warn "Temp directory not found: $TEMP_DIR"
    fi
    
    return 0
}

# 压缩备份
compress_backup() {
    local backup_path="$1"
    log_info "Compressing backup..."
    
    if [ "$BACKUP_COMPRESS" = "true" ]; then
        local compressed_name="${backup_path}.tar.gz"
        
        tar -czf "$compressed_name" -C "$(dirname "$backup_path")" "$(basename "$backup_path")"
        
        if [ -f "$compressed_name" ]; then
            # 删除原始目录
            rm -rf "$backup_path"
            
            backup_path="$compressed_name"
            log_info "Backup compressed: $backup_path"
            log_info "Compressed backup size: $(du -h "$backup_path" | cut -f1)"
        else
            log_error "Failed to compress backup"
            return 1
        fi
    fi
    
    echo "$backup_path"
}

# 加密备份
encrypt_backup() {
    local backup_path="$1"
    log_info "Encrypting backup..."
    
    if [ "$BACKUP_ENCRYPT" = "true" ] && [ -n "$ENCRYPTION_KEY" ]; then
        local encrypted_name="${backup_path}.enc"
        
        # 使用openssl加密
        openssl enc -aes-256-cbc -salt -in "$backup_path" -out "$encrypted_name" -k "$ENCRYPTION_KEY" || {
            log_error "Failed to encrypt backup"
            return 1
        }
        
        if [ -f "$encrypted_name" ]; then
            # 删除未加密的备份
            rm -f "$backup_path"
            
            backup_path="$encrypted_name"
            log_info "Backup encrypted: $backup_path"
            log_info "Encrypted backup size: $(du -h "$backup_path" | cut -f1)"
        else
            log_error "Failed to encrypt backup"
            return 1
        fi
    else
        log_info "Backup encryption skipped (disabled or no key provided)"
    fi
    
    echo "$backup_path"
}

# 创建备份元数据
create_backup_metadata() {
    local backup_path="$1"
    local metadata_file="${backup_path}.metadata"
    
    log_info "Creating backup metadata..."
    
    cat > "$metadata_file" << EOF
{
    "backup_created": "$(date -Iseconds)",
    "backup_path": "$backup_path",
    "backup_type": "full",
    "backup_compressed": $BACKUP_COMPRESS,
    "backup_encrypted": $BACKUP_ENCRYPT,
    "database_file": "$DATABASE_FILE",
    "data_dir": "$DATA_DIR",
    "log_dir": "$LOG_DIR",
    "config_dir": "$CONFIG_DIR",
    "temp_dir": "$TEMP_DIR",
    "redis_url": "$REDIS_URL",
    "system_info": {
        "hostname": "$(hostname)",
        "os": "$(uname -s)",
        "kernel": "$(uname -r)",
        "arch": "$(uname -m)"
    },
    "disk_usage": {
        "backup_size": "$(du -h "$backup_path" | cut -f1)",
        "data_dir_size": "$(du -sh "$DATA_DIR" 2>/dev/null | cut -f1 || echo 'N/A')",
        "log_dir_size": "$(du -sh "$LOG_DIR" 2>/dev/null | cut -f1 || echo 'N/A')"
    }
}
EOF
    
    log_info "Backup metadata created: $metadata_file"
}

# 验证备份
verify_backup() {
    local backup_path="$1"
    log_info "Verifying backup..."
    
    local backup_size=$(du -b "$backup_path" 2>/dev/null | cut -f1)
    
    if [ -z "$backup_size" ] || [ "$backup_size" -eq 0 ]; then
        log_error "Backup file is empty or doesn't exist"
        return 1
    fi
    
    log_info "Backup verification passed"
    log_info "Backup size: $backup_size bytes"
    
    return 0
}

# 清理旧备份
cleanup_old_backups() {
    log_info "Cleaning up old backups..."
    
    if [ -d "$BACKUP_DIR" ]; then
        # 删除超过保留期的备份
        find "$BACKUP_DIR" -name "backup_*" -type d -mtime +$BACKUP_RETENTION_DAYS -exec rm -rf {} \; 2>/dev/null || true
        find "$BACKUP_DIR" -name "backup_*.tar.gz" -mtime +$BACKUP_RETENTION_DAYS -delete 2>/dev/null || true
        find "$BACKUP_DIR" -name "backup_*.tar.gz.enc" -mtime +$BACKUP_RETENTION_DAYS -delete 2>/dev/null || true
        
        log_info "Old backups cleaned up (retention: $BACKUP_RETENTION_DAYS days)"
    fi
}

# 完整备份
full_backup() {
    log_info "Starting full backup..."
    
    # 创建备份目录
    backup_path=$(create_backup_dir)
    
    # 执行各项备份
    backup_database "$backup_path"
    backup_redis "$backup_path"
    backup_data "$backup_path"
    backup_logs "$backup_path"
    backup_config "$backup_path"
    backup_temp "$backup_path"
    
    # 创建备份清单
    cat > "$backup_path/manifest.txt" << EOF
Backup Manifest
===============
Created: $(date)
Backup Type: Full Backup
Backup Path: $backup_path

Contents:
EOF
    find "$backup_path" -type f -ls >> "$backup_path/manifest.txt"
    
    # 压缩备份
    backup_path=$(compress_backup "$backup_path")
    
    # 加密备份
    backup_path=$(encrypt_backup "$backup_path")
    
    # 创建元数据
    create_backup_metadata "$backup_path"
    
    # 验证备份
    verify_backup "$backup_path"
    
    # 清理旧备份
    cleanup_old_backups
    
    log_info "Full backup completed successfully: $backup_path"
    echo "$backup_path"
}

# 恢复数据库
restore_database() {
    local backup_path="$1"
    log_info "Restoring database..."
    
    # 解压备份（如果需要）
    local temp_dir="$TEMP_DIR/restore_$(date +%s)"
    mkdir -p "$temp_dir"
    
    if [[ "$backup_path" == *.tar.gz ]]; then
        tar -xzf "$backup_path" -C "$temp_dir"
        backup_path=$(find "$temp_dir" -name "backup_*" -type d | head -1)
    fi
    
    if [ -f "$backup_path/database/app.db" ]; then
        # 备份当前数据库
        if [ -f "$DATABASE_FILE" ]; then
            cp "$DATABASE_FILE" "$DATABASE_FILE.backup.$(date +%Y%m%d_%H%M%S)"
        fi
        
        # 恢复数据库
        cp "$backup_path/database/app.db" "$DATABASE_FILE"
        
        log_info "Database restored successfully"
    else
        log_error "Database backup not found"
        return 1
    fi
    
    # 清理临时目录
    rm -rf "$temp_dir"
    
    return 0
}

# 恢复数据
restore_data() {
    local backup_path="$1"
    log_info "Restoring data directory..."
    
    # 解压备份（如果需要）
    local temp_dir="$TEMP_DIR/restore_$(date +%s)"
    mkdir -p "$temp_dir"
    
    if [[ "$backup_path" == *.tar.gz ]]; then
        tar -xzf "$backup_path" -C "$temp_dir"
        backup_path=$(find "$temp_dir" -name "backup_*" -type d | head -1)
    fi
    
    if [ -f "$backup_path/data/data_files.tar" ]; then
        # 创建数据目录（如果不存在）
        mkdir -p "$DATA_DIR"
        
        # 恢复数据文件
        tar -xf "$backup_path/data/data_files.tar" -C "$DATA_DIR"
        
        log_info "Data directory restored successfully"
    else
        log_error "Data backup not found"
        return 1
    fi
    
    # 清理临时目录
    rm -rf "$temp_dir"
    
    return 0
}

# 恢复日志
restore_logs() {
    local backup_path="$1"
    log_info "Restoring logs directory..."
    
    # 解压备份（如果需要）
    local temp_dir="$TEMP_DIR/restore_$(date +%s)"
    mkdir -p "$temp_dir"
    
    if [[ "$backup_path" == *.tar.gz ]]; then
        tar -xzf "$backup_path" -C "$temp_dir"
        backup_path=$(find "$temp_dir" -name "backup_*" -type d | head -1)
    fi
    
    if [ -f "$backup_path/logs/logs.tar" ]; then
        # 创建日志目录（如果不存在）
        mkdir -p "$LOG_DIR"
        
        # 恢复日志文件
        tar -xf "$backup_path/logs/logs.tar" -C "$LOG_DIR"
        
        log_info "Logs directory restored successfully"
    else
        log_error "Logs backup not found"
        return 1
    fi
    
    # 清理临时目录
    rm -rf "$temp_dir"
    
    return 0
}

# 完整恢复
full_restore() {
    local backup_path="$1"
    log_info "Starting full restore from: $backup_path"
    
    # 检查备份文件是否存在
    if [ ! -f "$backup_path" ]; then
        log_error "Backup file not found: $backup_path"
        exit 1
    fi
    
    # 停止应用服务（如果运行）
    if pgrep -f "uvicorn.*app.main:app" > /dev/null; then
        log_info "Stopping application service..."
        pkill -f "uvicorn.*app.main:app" || true
        sleep 5
    fi
    
    # 执行各项恢复
    restore_database "$backup_path"
    restore_data "$backup_path"
    restore_logs "$backup_path"
    
    # 恢复配置文件（可选）
    if [ "$2" = "--with-config" ]; then
        log_info "Restoring configuration files..."
        # 这里添加配置恢复逻辑
    fi
    
    log_info "Full restore completed successfully"
    log_info "Please restart the application service"
}

# 列出备份
list_backups() {
    log_info "Listing available backups..."
    
    if [ ! -d "$BACKUP_DIR" ]; then
        log_warn "Backup directory not found: $BACKUP_DIR"
        return 1
    fi
    
    echo "Available Backups:"
    echo "=================="
    
    find "$BACKUP_DIR" -name "backup_*" -type d -exec basename {} \; | sort -r
    find "$BACKUP_DIR" -name "backup_*.tar.gz" -exec basename {} \; | sort -r
    find "$BACKUP_DIR" -name "backup_*.tar.gz.enc" -exec basename {} \; | sort -r
    
    echo ""
    echo "Backup directory: $BACKUP_DIR"
    echo "Total backups: $(find "$BACKUP_DIR" -name "backup_*" | wc -l)"
}

# 显示帮助信息
show_help() {
    echo "Backup and Restore Script"
    echo ""
    echo "Usage: $0 [COMMAND] [OPTIONS]"
    echo ""
    echo "Commands:"
    echo "  backup                     Perform full backup"
    echo "  restore <backup_file>      Restore from backup"
    echo "  restore <backup_file> --with-config  Restore including configuration"
    echo "  list                       List available backups"
    echo "  cleanup                    Clean up old backups"
    echo "  help                       Show this help message"
    echo ""
    echo "Environment Variables:"
    echo "  BACKUP_DIR                 Backup directory (default: ./backups)"
    echo "  DATA_DIR                   Data directory (default: ./data)"
    echo "  LOG_DIR                    Log directory (default: ./logs)"
    echo "  BACKUP_RETENTION_DAYS      Backup retention in days (default: 30)"
    echo "  BACKUP_COMPRESS            Compress backups (default: true)"
    echo "  BACKUP_ENCRYPT              Encrypt backups (default: false)"
    echo "  ENCRYPTION_KEY             Encryption key for backups"
    echo "  DATABASE_FILE              Database file path"
    echo "  REDIS_URL                  Redis connection URL"
    echo ""
    echo "Examples:"
    echo "  $0 backup                                          # Perform full backup"
    echo "  $0 restore ./backups/backup_20240101_120000.tar.gz  # Restore from backup"
    echo "  $0 list                                            # List available backups"
    echo "  BACKUP_COMPRESS=false $0 backup                  # Backup without compression"
}

# 主函数
main() {
    local command="$1"
    local backup_file="$2"
    
    # 创建必要目录
    mkdir -p "$BACKUP_DIR" "$DATA_DIR" "$LOG_DIR" "$TEMP_DIR"
    
    # 检查必要的命令
    check_command tar
    check_command sqlite3
    check_command du
    check_command find
    
    case $command in
        "backup")
            full_backup
            ;;
        "restore")
            if [ -z "$backup_file" ]; then
                log_error "Please specify backup file to restore from"
                show_help
                exit 1
            fi
            full_restore "$backup_file" "$3"
            ;;
        "list")
            list_backups
            ;;
        "cleanup")
            cleanup_old_backups
            ;;
        "help"|"--help"|"-h")
            show_help
            ;;
        "")
            log_error "No command specified"
            show_help
            exit 1
            ;;
        *)
            log_error "Unknown command: $command"
            show_help
            exit 1
            ;;
    esac
}

# 运行主函数
main "$@"
